import ImageUtils
import shutil
import os

class ClassSet:
    def __init__(self, path):
        self.rootPath = path  # 根路径
        self.trainPath = path + "/train"
        self.testPath = path + "/test"
        # 检查创建目录
        ImageUtils.check_and_create_dir(self.trainPath)
        ImageUtils.check_and_create_dir(self.testPath)
        self.initDate()

    def initDate(self):

        # 获取默认的初始图片
        self.originImageList = ImageUtils.getFilePathFromDir(self.rootPath, [".jpg", ".png",".jpeg"])  # 未分配的照片集合
        self.trainImageList = ImageUtils.getFilePathFromDir(self.trainPath, [".png", ".jpg",".jpeg"])  # 训练的照片集合
        self.testImageList = ImageUtils.getFilePathFromDir(self.testPath, [".png", ".jpg",".jpeg"])  # 测试照片集合
        print("未去重未分配的原始图片" + str(len(self.originImageList)) + "个" + str(self.originImageList))
        print("未去重未分配的训练图片" + str(len(self.trainImageList)) + "个" + str(self.trainImageList))
        print("未去重未分配的测试图片" + str(len(self.testImageList)) + "个" + str(self.testImageList))
        unionTestAndTrainSet = list(map(ImageUtils.getFileNameNotSufFromPath, list(set(self.trainImageList) | set(self.testImageList))))
        print("并集后的图片"+str(unionTestAndTrainSet))
        repeat_item = []
        for item in self.originImageList:
            if ImageUtils.getFileNameNotSufFromPath(item) in unionTestAndTrainSet:
                repeat_item.append(item)
        print("重复的图片"+ str(len(repeat_item)) + "个" + str(repeat_item))
        self.originImageList = list(set(self.originImageList) - set(repeat_item))
        print("去重后未分配的图片"+ str(len(self.originImageList)) + "个" + str(list(self.originImageList)))
        for item in repeat_item:
            xmlPath = item.split(".")[0] + ".xml"
            if os.path.exists(xmlPath):
                os.remove(xmlPath)
            os.remove(item)
            print("删除重复图片"+item)

        # originList = map(ImageUtils.getFileNameNotSufFromPath, self.originImageList)
        # print("等待去重经过获取文件名过滤的并集图片" + str(list(unionTestAndTrainSet)))
        # print("等待去重经过获取文件名过滤的原始图片" + str(list(originList)))
        # repeat_item = set(originList) & set(unionTestAndTrainSet)

    def changeImageFile(self, new_trainList, new_testList, new_origin):
        """
        根据新的列表改变文件位置
        :param new_trainList:移动后的训练图片
        :param new_testList: 移动后的测试图片
        :param new_origin:   移动或的原始图片
        :return:
        """
        # 将新增的照片从源位置移入新位置
        for train in new_trainList:
            if train not in self.trainImageList:
                xmlPath = train.split(".")[0] + ".xml"
                shutil.move(train, self.trainPath)
                shutil.move(xmlPath, self.trainPath)

        for test in new_testList:
            if test not in self.testImageList:
                xmlPath = test.split(".")[0] + ".xml"
                shutil.move(test, self.testPath)
                shutil.move(xmlPath, self.testPath)

        for origin in new_origin:
            if origin not in self.originImageList:
                xmlPath = origin.split(".")[0] + ".xml"
                shutil.move(origin, self.rootPath)
                shutil.move(xmlPath, self.rootPath)

    def saveClass(self):
        # 删除原来的重新复制一份
        ImageUtils.removeAllFileFromDir(self.trainPath, [".jpg", ".png",".jpeg", ".xml"])
        ImageUtils.removeAllFileFromDir(self.testPath, [".jpg", ".png",".jpeg", ".xml"])
        for item in self.trainImageList:
            xmlPath = item.split(".")[0] + ".xml"
            shutil.copy(item, self.trainPath)
            shutil.copy(xmlPath, self.trainPath)
        for item in self.testImageList:
            xmlPath = item.split(".")[0] + ".xml"
            shutil.copy(item, self.testPath)
            shutil.copy(xmlPath, self.testPath)
